#include <iostream>
#include <cuda_runtime.h>
#include <chrono>

// 打印 CUDA 错误信息
void checkCudaError(cudaError_t err, const char* msg) {
    if (err != cudaSuccess) {
        std::cerr << msg << ": " << cudaGetErrorString(err) << std::endl;
        exit(EXIT_FAILURE);
    }
}

// 测试主机到设备的数据拷贝带宽
void testHostToDeviceBandwidth(size_t dataSize) {
    float* hostData;
    float* deviceData;

    // 分配主机内存
    hostData = new float[dataSize / sizeof(float)];
    // 初始化主机数据
    for (size_t i = 0; i < dataSize / sizeof(float); ++i) {
        hostData[i] = static_cast<float>(i);
    }

    // 分配设备内存
    checkCudaError(cudaMalloc((void**)&deviceData, dataSize), "cudaMalloc failed");

    // 预热
    checkCudaError(cudaMemcpy(deviceData, hostData, dataSize, cudaMemcpyHostToDevice), "cudaMemcpy (warmup) failed");

    // 开始计时
    auto start = std::chrono::high_resolution_clock::now();

    // 进行多次数据拷贝以减少计时误差
    const int numIterations = 100;
    for (int i = 0; i < numIterations; ++i) {
        checkCudaError(cudaMemcpy(deviceData, hostData, dataSize, cudaMemcpyHostToDevice), "cudaMemcpy failed");
    }

    // 同步设备以确保所有拷贝操作完成
    checkCudaError(cudaDeviceSynchronize(), "cudaDeviceSynchronize failed");

    // 结束计时
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed = end - start;

    // 计算带宽
    double bandwidth = (static_cast<double>(dataSize * numIterations) / (1024.0 * 1024.0 * 1024.0)) / elapsed.count();

    std::cout << "Host to Device Bandwidth: " << bandwidth << " GB/s" << std::endl;

    // 释放内存
    delete[] hostData;
    checkCudaError(cudaFree(deviceData), "cudaFree failed");
}

// 测试设备到主机的数据拷贝带宽
void testDeviceToHostBandwidth(size_t dataSize) {
    float* hostData;
    float* deviceData;

    // 分配主机内存
    hostData = new float[dataSize / sizeof(float)];

    // 分配设备内存
    checkCudaError(cudaMalloc((void**)&deviceData, dataSize), "cudaMalloc failed");

    // 初始化设备数据
    checkCudaError(cudaMemset(deviceData, 0, dataSize), "cudaMemset failed");

    // 预热
    checkCudaError(cudaMemcpy(hostData, deviceData, dataSize, cudaMemcpyDeviceToHost), "cudaMemcpy (warmup) failed");

    // 开始计时
    auto start = std::chrono::high_resolution_clock::now();

    // 进行多次数据拷贝以减少计时误差
    const int numIterations = 100;
    for (int i = 0; i < numIterations; ++i) {
        checkCudaError(cudaMemcpy(hostData, deviceData, dataSize, cudaMemcpyDeviceToHost), "cudaMemcpy failed");
    }

    // 同步设备以确保所有拷贝操作完成
    checkCudaError(cudaDeviceSynchronize(), "cudaDeviceSynchronize failed");

    // 结束计时
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed = end - start;

    // 计算带宽
    double bandwidth = (static_cast<double>(dataSize * numIterations) / (1024.0 * 1024.0 * 1024.0)) / elapsed.count();

    std::cout << "Device to Host Bandwidth: " << bandwidth << " GB/s" << std::endl;

    // 释放内存
    delete[] hostData;
    checkCudaError(cudaFree(deviceData), "cudaFree failed");
}

int main() {
    // 数据大小，例如 1024 MB
    size_t dataSize = 1024 * 1024 * 1024;

    // 测试主机到设备的带宽
    testHostToDeviceBandwidth(dataSize);

    // 测试设备到主机的带宽
    testDeviceToHostBandwidth(dataSize);

    return 0;
}    